import pymysql
import psycopg2
import psycopg2.extras
from psycopg2.extras import RealDictCursor
import pandas as pd
from psycopg2 import sql
from psycopg2.extras import DictCursor
from typing import Union
import requests
import pandas as pd

class DatabaseService:
    """数据库连接和操作服务"""
    
    @staticmethod
    def test_connection(db_type, host, port, username, password, database):
        """测试数据库连接"""
        try:
            if db_type == 'mysql':
                connection = pymysql.connect(
                    host=host,
                    port=port,
                    user=username,
                    password=password,
                    database=database
                )
                connection.close()
            elif db_type == 'postgresql':
                connection = psycopg2.connect(
                    host=host,
                    port=port,
                    user=username,
                    password=password,
                    dbname=database
                )
                connection.close()
            else:
                raise ValueError(f"不支持的数据库类型: {db_type}")
            
            return True, "连接成功"
        except Exception as e:
            return False, str(e)
    
    @staticmethod
    def get_connection(connection_config):
        """获取数据库连接"""
        if connection_config['db_type'] == 'mysql':
            return pymysql.connect(
                host=connection_config['host'],
                port=connection_config['port'],
                user=connection_config['user'],
                password=connection_config['password'],
                database=connection_config['database'],
                cursorclass=pymysql.cursors.DictCursor
            )
        elif connection_config['db_type'] == 'postgresql' or connection_config['db_type'] == 'openGauss':
            return psycopg2.connect(
                host=connection_config['host'],
                port=connection_config['port'],
                user=connection_config['user'],
                password=connection_config['password'],
                dbname=connection_config['database'],
                cursor_factory=RealDictCursor
            )
        else:
            raise ValueError(f"Unsupported database type: {connection_config['db_type']}")
    
    @staticmethod
    def execute_query(connection_config, query, limit=100):
        """执行SQL查询并返回结果"""
        connection = None
        try:
            connection = DatabaseService.get_connection(connection_config)
            if limit:
                wrapped_query = f"SELECT * FROM ({query}) as temp_table LIMIT {limit}"
            else:
                wrapped_query = query
            df = pd.read_sql(wrapped_query, connection)
            print("Query executed successfully: ", df.shape)
            return True, {
                "columns": df.columns.tolist(),
                "data": df.to_dict('records'),
                "total_rows": len(df)
            }
        except Exception as e:
            print("Query execution failed: ", str(e))
            return False, str(e)
        finally:
            if connection:
                connection.close()

    @staticmethod
    def run_sql_pg(sql_query: str, db_config: dict) -> Union[pd.DataFrame, None]:
        try:
            with psycopg2.connect(
                dbname=db_config["dbname"], 
                user=db_config["user"], 
                password=db_config["password"], 
                host=db_config["host"], 
                port=db_config["port"]
            ) as conn:
                with conn.cursor(cursor_factory=DictCursor) as cursor:
                    cursor.execute(sql_query)
                    results = cursor.fetchall()
                    
                    df = pd.DataFrame(results, columns=[desc[0] for desc in cursor.description])
                    return df
        except psycopg2.Error as e:
            raise Exception(f"Database error: {e}")
        except Exception as e:
            raise e
    
    @staticmethod
    def run_sql_iris(sql_query: str, db_config: dict) -> Union[pd.DataFrame, None]:
        try:
            # Ensure the required dependency is installed
            import sqlalchemy_iris  # noqa: F401
        except ImportError:
            raise Exception(
                "You need to install required dependencies to execute this method, "
                "run command: \npip install sqlalchemy-iris"
            )

        # Retrieve configuration values from db_config dictionary
        host = db_config.get("host")
        if not host:
            raise Exception("Please set your IRIS host")

        dbname = db_config.get("dbname")
        if not dbname:
            raise Exception("Please set your IRIS database")

        user = db_config.get("user")
        if not user:
            raise Exception("Please set your IRIS user")

        password = db_config.get("password")
        if not password:
            raise Exception("Please set your IRIS password")
        # Ensure the password is URL-encoded
        from urllib.parse import quote_plus
        password = quote_plus(password)

        port = db_config.get("port")
        if not port:
            raise Exception("Please set your IRIS port")

        # Create the SQLAlchemy engine for IRIS
        from sqlalchemy import create_engine, text
        from sqlalchemy.exc import SQLAlchemyError

        try:
            engine = create_engine(f"iris://{user}:{password}@{host}:{port}/{dbname}")
        except SQLAlchemyError as e:
            raise Exception(f"Error creating engine: {e}")

        # Execute the SQL query and return the results as a DataFrame
        try:
            with engine.connect() as connection:
                result = connection.execute(text(sql_query))
                results = result.fetchall()
                df = pd.DataFrame(results, columns=result.keys())
                return df
        except SQLAlchemyError as e:
            raise Exception(f"Database error: {e}")
        except Exception as e:
            raise e

    @staticmethod
    def run_sql_kylin(sql_query: str, db_config: dict) -> Union[pd.DataFrame, None]:
        """
        连接 Apache Kylin 执行 SQL 查询，并返回结果 DataFrame。
        
        参数:
            sql_query: 要执行的 SQL 查询语句。
            db_config: 包含连接参数的字典，至少需包含以下键：
                - host: Kylin 服务的主机地址
                - port: Kylin 服务的端口号
                - user: 用户名
                - password: 密码
                - project: 要查询的 Kylin 项目名称（默认为 default）
                
        返回:
            查询结果对应的 pandas DataFrame；如果没有结果则返回 None。
        """
        try:
            # 拼接请求 URL
            url = f"http://{db_config['host']}:{db_config['port']}/kylin/api/query"
            
            # 构造请求头和认证信息
            headers = {"Content-Type": "application/json"}
            auth = (db_config["user"], db_config["password"])
            
            # 构造请求负载，注意项目名称可从 db_config 中获取
            payload = {
                "sql": sql_query,
                "project": db_config.get("project", "default")
            }
            
            # 发送 POST 请求
            response = requests.post(url, json=payload, headers=headers, auth=auth)
            response.raise_for_status()  # 若请求失败则抛出异常
            
            # 解析返回 JSON 数据
            result = response.json()
            
            # 从返回结果中提取字段信息和数据行
            columns = [col["label"] for col in result.get("columnMetas", [])]
            rows = result.get("results", [])
            df = DatabaseService.convert_df_types(pd.DataFrame(rows, columns=columns), result.get("columnMetas", []))
            return df
        except requests.exceptions.RequestException as e:
            raise Exception(f"HTTP 请求错误: {e}")
        except Exception as e:
            raise e
    
    @staticmethod
    def run_sql_sqlserver(sql_query: str, db_config: dict) -> Union[pd.DataFrame, None]:
        """
        连接 SQL Server 执行 SQL 查询，并返回结果 DataFrame。
        需要安装 sqlalchemy 和 pyodbc: pip install sqlalchemy pyodbc
        """
        try:
            from sqlalchemy import create_engine, text
            user = db_config["user"]
            password = db_config["password"]
            host = db_config["host"]
            port = db_config["port"]
            database = db_config["database"]
            engine = create_engine(f"mssql+pyodbc://{user}:{password}@{host}:{port}/{database}?driver=ODBC+Driver+17+for+SQL+Server")
            with engine.connect() as connection:
                print("sql_query: ", sql_query)
                result = connection.execute(text(sql_query))
                results = result.fetchall()
                df = pd.DataFrame(results, columns=result.keys())
                return df
        except Exception as e:
            raise Exception(f"SQLServer error: {e}")

    @staticmethod
    def run_sql_auto(sql_query: str, db_config: dict) -> Union[pd.DataFrame, None]:
        """
        根据 db_config['db_type'] 自动选择查询方式。
        支持: kylin, iris, sqlserver
        """
        db_type = db_config.get("db_type", "").lower()
        if db_type == "kylin":
            return DatabaseService.run_sql_kylin(sql_query, db_config)
        elif db_type == "iris":
            return DatabaseService.run_sql_iris(sql_query, db_config)
        elif db_type == "sqlserver":
            return DatabaseService.run_sql_sqlserver(sql_query, db_config)
        else:
            raise Exception(f"不支持的数据库类型: {db_type}")

    @staticmethod
    def count_query_results(connection_config, query):
        """计算查询结果的总行数"""
        connection = None
        try:
            connection = DatabaseService.get_connection(connection_config)
            count_query = f"SELECT COUNT(*) as total_count FROM ({query}) as temp_table"
            result = pd.read_sql(count_query, connection)
            total_count = result['total_count'].iloc[0]
            return True, total_count
        except Exception as e:
            return False, str(e)
        finally:
            if connection:
                connection.close()
    
    @staticmethod
    def stream_query_results(connection_config, query, batch_size=1000):
        """流式获取查询结果"""
        connection = None
        cursor = None
        try:
            connection = DatabaseService.get_connection(connection_config)
            if connection_config.db_type == 'mysql':
                cursor = connection.cursor(pymysql.cursors.DictCursor)
            elif connection_config.db_type == 'postgresql':
                cursor = connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
            cursor.execute(query)
            while True:
                rows = cursor.fetchmany(batch_size)
                if not rows:
                    break
                yield rows
        except Exception as e:
            raise Exception(f"Query execution failed: {str(e)}")
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    @staticmethod
    def convert_df_types(df, columnMetas):
        """
        根据 columnMetas 中的字段类型转换 df 中的列类型。
        :param df: pandas DataFrame
        :param columnMetas: list of dict，每个 dict 包含 'columnName' 和 'dataType' 等信息
        :return: 转换类型后的 DataFrame
        """
        # 定义一个映射关系：Kylin 类型 -> pandas dtype
        type_mapping = {
            "double": "float64",
            "float": "float64",
            "bigint": "int64",
            "integer": "int64",
            "int": "int64",
            "decimal": "float64",
            "timestamp": "datetime64[ns]",
            "date": "datetime64[ns]",
            # 对于字符串类型或者其他不需要转换的类型，可以设置为 None
            "string": None,
            "varchar": None,
            # 可根据需要添加其他映射
        }
        
        for meta in columnMetas:
            col = meta.get("label")
            col_type = meta.get("columnTypeName", "").lower()
            target_dtype = type_mapping.get(col_type)
            
            # 如果未指定转换或者列名不在 df 中，则跳过
            if target_dtype is None or col not in df.columns:
                continue
            
            # 对于 datetime 类型，使用 pd.to_datetime
            if target_dtype.startswith("datetime"):
                try:
                    df[col] = pd.to_datetime(df[col], errors="coerce")
                except Exception as e:
                    print(f"转换列 {col} 为 datetime 类型时出错: {e}")
            else:
                try:
                    df[col] = df[col].astype(target_dtype)
                except Exception as e:
                    print(f"转换列 {col} 为 {target_dtype} 时出错: {e}")
        
        return df